#!/usr/bin/env python3
"""
数据库初始化脚本

该脚本用于：
1. 连接PostgreSQL服务器
2. 创建项目数据库
3. 创建应用用户并设置权限
4. 创建数据库表结构
"""

import asyncio
import os
import sys
from typing import Optional

import asyncpg
from sqlalchemy import create_engine, text
from sqlalchemy.ext.asyncio import create_async_engine

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.core.database import Base
from src.core.logger import get_logger

logger = get_logger(__name__)

# 数据库配置
DB_CONFIG = {
    'host': '101.35.56.140',
    'port': 5432,
    'admin_user': 'postgres', 
    'admin_password': 'psl#@!logIN4SS2000',
    'app_user': '_all_topic_rw',
    'app_password': 'ShowAllTopicRW!!!2025',
    'database': 'show_all_topic',
    'test_database': 'show_all_topic_test'
}


async def test_connection(host: str, port: int, user: str, password: str, database: str = None) -> bool:
    """测试数据库连接"""
    try:
        if database:
            conn = await asyncpg.connect(
                host=host, port=port, user=user, password=password, database=database
            )
        else:
            conn = await asyncpg.connect(
                host=host, port=port, user=user, password=password
            )
        
        version = await conn.fetchval('SELECT version()')
        logger.info(f"连接成功: {user}@{host}:{port}")
        logger.info(f"PostgreSQL版本: {version}")
        await conn.close()
        return True
        
    except Exception as e:
        logger.error(f"连接失败: {user}@{host}:{port}/{database or 'default'}")
        logger.error(f"错误信息: {e}")
        return False


async def create_database_and_user():
    """创建数据库和用户"""
    config = DB_CONFIG
    
    # 1. 测试管理员连接
    logger.info("=== 测试管理员连接 ===")
    admin_connected = await test_connection(
        config['host'], config['port'], 
        config['admin_user'], config['admin_password']
    )
    
    if not admin_connected:
        logger.error("无法连接管理员账号，请检查用户名和密码")
        return False
    
    # 2. 连接到postgres数据库执行管理操作
    try:
        conn = await asyncpg.connect(
            host=config['host'], port=config['port'],
            user=config['admin_user'], password=config['admin_password'],
            database='postgres'  # 连接到默认的postgres数据库
        )
        
        # 3. 检查数据库是否存在
        db_exists = await conn.fetchval(
            "SELECT 1 FROM pg_database WHERE datname = $1", 
            config['database']
        )
        
        if not db_exists:
            logger.info(f"创建数据库: {config['database']}")
            await conn.execute(f'CREATE DATABASE "{config['database']}"')
        else:
            logger.info(f"数据库已存在: {config['database']}")
        
        # 创建测试数据库
        test_db_exists = await conn.fetchval(
            "SELECT 1 FROM pg_database WHERE datname = $1", 
            config['test_database']
        )
        
        if not test_db_exists:
            logger.info(f"创建测试数据库: {config['test_database']}")
            await conn.execute(f'CREATE DATABASE "{config['test_database']}"')
        else:
            logger.info(f"测试数据库已存在: {config['test_database']}")
        
        # 4. 检查用户是否存在
        user_exists = await conn.fetchval(
            "SELECT 1 FROM pg_user WHERE usename = $1", 
            config['app_user']
        )
        
        if not user_exists:
            logger.info(f"创建用户: {config['app_user']}")
            await conn.execute(f"""
                CREATE USER "{config['app_user']}" WITH PASSWORD '{config['app_password']}'
            """)
        else:
            logger.info(f"用户已存在: {config['app_user']}")
            # 更新密码
            await conn.execute(f"""
                ALTER USER "{config['app_user']}" WITH PASSWORD '{config['app_password']}'
            """)
        
        # 5. 设置用户权限
        logger.info("设置用户权限...")
        await conn.execute(f'GRANT ALL PRIVILEGES ON DATABASE "{config['database']}" TO "{config['app_user']}"')
        await conn.execute(f'GRANT ALL PRIVILEGES ON DATABASE "{config['test_database']}" TO "{config['app_user']}"')
        
        # 6. 设置schema权限
        logger.info("设置schema权限...")
        for db_name in [config['database'], config['test_database']]:
            await conn.execute(f'GRANT CONNECT ON DATABASE "{db_name}" TO "{config['app_user']}"')
            
        await conn.close()
        
        # 连接到各个数据库设置schema权限
        for db_name in [config['database'], config['test_database']]:
            db_conn = await asyncpg.connect(
                host=config['host'], port=config['port'],
                user=config['admin_user'], password=config['admin_password'],
                database=db_name
            )
            
            await db_conn.execute(f'GRANT ALL ON SCHEMA public TO "{config['app_user']}"')
            await db_conn.execute(f'GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO "{config['app_user']}"')
            await db_conn.execute(f'GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public TO "{config['app_user']}"')
            await db_conn.execute(f'ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO "{config['app_user']}"')
            await db_conn.execute(f'ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO "{config['app_user']}"')
            
            await db_conn.close()
            logger.info(f"数据库 {db_name} 权限设置完成")
        
        logger.info("数据库和用户创建完成")
        return True
        
    except Exception as e:
        logger.error(f"创建数据库和用户失败: {e}")
        return False


async def create_tables():
    """创建数据库表"""
    config = DB_CONFIG
    
    # 测试应用用户连接
    logger.info("=== 测试应用用户连接 ===")
    app_connected = await test_connection(
        config['host'], config['port'],
        config['app_user'], config['app_password'],
        config['database']
    )
    
    if not app_connected:
        logger.error("无法连接应用用户账号")
        return False
    
    # 创建数据库表
    try:
        database_url = (
            f"postgresql+asyncpg://{config['app_user']}:{config['app_password']}"
            f"@{config['host']}:{config['port']}/{config['database']}"
        )
        
        logger.info("创建数据库表...")
        engine = create_async_engine(database_url)
        
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        
        await engine.dispose()
        logger.info("数据库表创建完成")
        
        # 为测试数据库也创建表
        test_database_url = (
            f"postgresql+asyncpg://{config['app_user']}:{config['app_password']}"
            f"@{config['host']}:{config['port']}/{config['test_database']}"
        )
        
        logger.info("创建测试数据库表...")
        test_engine = create_async_engine(test_database_url)
        
        async with test_engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        
        await test_engine.dispose()
        logger.info("测试数据库表创建完成")
        
        return True
        
    except Exception as e:
        logger.error(f"创建数据库表失败: {e}")
        return False


async def main():
    """主函数"""
    logger.info("开始数据库初始化...")
    
    # 1. 创建数据库和用户
    if not await create_database_and_user():
        logger.error("数据库和用户创建失败")
        return False
    
    # 2. 创建数据库表
    if not await create_tables():
        logger.error("数据库表创建失败")
        return False
    
    logger.info("数据库初始化完成！")
    return True


if __name__ == "__main__":
    success = asyncio.run(main())
    if not success:
        sys.exit(1)